from PIL import Image
import os
import numpy as np
import torch
import math
import cv2
import matplotlib.pyplot as plt

from net.test_ import test_2

def grasp_box():
    blend_img, mask = test_2()
    # 找mask的轮廓，将轮廓上的坐标存储到contours
    imgray=cv2.cvtColor(mask,cv2.COLOR_BGR2GRAY)
    ret, binary = cv2.threshold(imgray,127,255,cv2.THRESH_BINARY) 
    # contours是一个列表
    img, contours, __ = cv2.findContours(binary,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
    if len(contours) > 0:
        c = max(contours, key=cv2.contourArea)
    
    # convert the PIL Image into a numpy array
    mask = np.array(mask)

    # 取颜色不同的实例对象
    obj_ids = np.unique(mask)
    # first id is the background, so remove it
    obj_ids = obj_ids[1:]
    # split the color-encoded mask into a set
    # of binary masks
    masks = mask == obj_ids[:, None, None]
    num_objs = len(obj_ids)
    for i in range(num_objs):
        pos = np.where(masks[i])
        # pos[1]、pos[0]包含mask上的全部x坐标、y坐标
        # 通过获取mask上的最大和最小的x、y坐标，来获取bbox
        xmin = np.min(pos[1])
        xmax = np.max(pos[1])
        ymin = np.min(pos[0])
        ymax = np.max(pos[0])

    # 获取可抓取区域的中心点坐标center
    x_center = (xmin + xmax)/2
    y_center = (ymin + ymax)/2  
    # 思路：先以center为中心，构建长轴和短轴的坐标系
    # 1、得到长轴与mask轮廓相交的两个点的坐标：
    #   得到xmin/xmax,ymin/ymax,x与y两两组合得到4组坐标，4组坐标两两求长度，取最大值对应的两个坐标
    # 2、计算过中心点，且与长轴垂直的两个点的坐标   
    
    # coord = np.array([[xmin,ymin],[xmin,ymax],[xmax,ymin],[xmax,ymax]])
    # max = []
    # for i in (len(coord)-1):
    #     for j in (len(coord)-1-i):
    #         value = np.square(coord[i][0] - coord[i+j+1][0]) + np.square(coord[i][1] - coord[i+j+1][1])
    #         max.append(value)
    
    max_ = []
    max = np.array(max_)
    for i in (len(pos[1])):
        for j in (len(pos[0])):
            value = np.square(pos[1][i]-pos[1][j]) + np.square(pos[0][i]-pos[0][j])
            max[i,j] = value
    
    i_, j_ = np.where(np.max(max))
    
    # 得到mask上长轴对应的两点
    x1 = pos[1][i_[0]]
    y1 = pos[0][i_[0]]

    x2 = pos[1][j_[0]]
    y2 = pos[0][j_[0]]

    # 计算长轴与x轴之间的角度: 直接计算长轴的斜率
    if y2 >=y1:
        k1 = (y2 - y1)/(x2 - x1)
        theta = math.atan(k1)

    else:
        k1 = (y1 - y2)/(x1 - x2)
        theta = math.atan(k1)

     

    # atan对应角度在-90 ~ 90
    theta = math.degrees(theta)
    if theta == 90:
        alpha = -90
    elif (theta < 90) and (theta > 0):
        alpha = theta -90
    elif (theta > -90) and (theta <= 0):
        alpha = theta
    # 将theta值转换为在0 ~180之间
    if theta < 0:
        theta += 180

    # 计算短轴的长度：
    # 判断mask轮廓所有点中满足arctan|(k2 - k1)/(1 + k1*k2)|==90的两个点
    store = []
    for k in len(c):
        k2 = (c[k][1] - y_center)/(c[k][0] - x_center)
        k3 = abs((k2-k1)/(1+k1*k2))
        # 可能会报错，atan不能直接处理分母为0的情况
        angle1 = math.atan(k3)
        if angle1 == math.radians(90):
            store.append(k)
    
    if len(store) == 2:
        idx1 = store[0]
        idx2 = store[1]

    min_blob = math.square(c[idx1][1] - c[idx2][1]) + math.square(c[idx1][0] - c[idx2][0])
    min_blob = math.sqrt(min_blob)
    
    rx = x_center
    ry = y_center
    w = 40 + min_blob
    h = 30
    R = ((rx, ry), (w, h), theta)
    # 把边框值赋给boxpoints()得到四个顶点
    rect = ((rx, ry), (w, h), alpha)
    points = cv2.boxPoints(rect)
    points = np.int0(points)
    
    # 按点画框: 2表示画线的宽度
    final_img=cv2.drawContours(blend_img,points,-1,(255,0,0),2)
    plt.imshow(final_img)
    plt.savefig('../net_structure_test/test_final.jpg')

    return R
        


        
         

                

            



